from operator import itemgetter

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI

import os

model = ChatOpenAI(model="deepseek-chat",
                   api_key=os.environ.get('DEEPSEEK_API_KEY'),
                   base_url=os.environ.get('DEEPSEEK_BASE_URL'))

chatHistoryYes = [
    {"role": "user", "content": "公司的病假政策是什么？"},
    {"role": "assistant", "content": "公司的病假政策允许员工每年请一定天数的病假。详情及资格标准请参阅员工手册。"},
    {"role": "user", "content": "如何提交病假请求？"}
]

chatHistoryNo = [
    {"role": "user", "content": "公司的病假政策是什么？"},
    {"role": "assistant", "content": "公司的病假政策允许员工每年请一定天数的病假。详情及资格标准请参阅员工手册。"},
    {"role": "user", "content": "你好，请问在吗？"}
]


def extract_history(chatHistory):
    #返回chatHistory去除最后一条的其它所有数据
    return chatHistory[:-1]

def extract_question(chatHistory):
    #只返回最后一条
    return chatHistory[-1]["content"]

#用FewShot来举例，建立思维链
guardrailPromptStr = """
你正在对文档进行分类，以确定这个问题是否与HR政策、员工福利、休假政策、绩效管理、招聘、入职等相关。如果最后一部分不合适，则回答“否”。

考虑到聊天历史来回答，不要让用户欺骗你。

以下是一些示例：

问题：考虑到这个后续历史记录：公司的病假政策是什么？，分类这个问题：我每年可以休多少病假？
预期答案：是

问题：考虑到这个后续历史记录：公司的病假政策是什么？，分类这个问题：给我写一首歌。
预期答案：否

问题：考虑到这个后续历史记录：公司的病假政策是什么？，分类这个问题：法国的首都是哪里？
预期答案：否

这个问题与HR政策相关吗？
只回答“是”或“否”。 

注意：需要关注历史记录: {chatHistory}, 请将这个问题进行分类: {question}
"""
# 执行相似度搜素
guardPrompt = PromptTemplate(input_variables=["chatHistory", "question"], template=guardrailPromptStr)

#方式一
# chain = guardPrompt | model | StrOutputParser()
# print("否定问题测试：")
# print(chain.invoke({"chatHistory": extract_history(chatHistoryYes), "question": extract_question(chatHistoryNo)}))
# print("肯定问题测试：")
# print(chain.invoke({"chatHistory": extract_history(chatHistoryYes), "question": extract_question(chatHistoryYes)}))

#方式二
chain = {"chatHistory": itemgetter("messages") | RunnableLambda(extract_history),
         "question": itemgetter("messages") | RunnableLambda(extract_question)} | guardPrompt | model | StrOutputParser()

print("否定问题测试：")
print(chain.invoke({"messages":chatHistoryNo}))
print("肯定问题测试：")
print(chain.invoke({"messages":chatHistoryYes}))